import os
import paramiko
from paramiko import SSHClient, AutoAddPolicy
import paramiko
from scp import SCPClient
from logger import Logger


class SSHTarget:
    """SSH target processing class"""

    def __init__(self, destination, ssh_key_file=None):
        """
        Initialize SSH target

        Parameters:
            destination (str): SSH target address, format: user@host:/path or ssh://user@host:/path
        """
        self.destination = destination
        self.username = None
        self.hostname = None
        self.remote_path = None
        self.password = None
        self.ssh_client = None
        self.sftp = None
        self.scp_client = None
        self.ssh_key_file = ssh_key_file

        self._parse_destination()

    def _parse_destination(self):
        """Parse SSH target address"""
        dest = self.destination

        # Handle ssh:// prefix
        if dest.startswith("ssh://"):
            dest = dest[6:]  # Remove 'ssh://' prefix

        if ":" not in dest:
            raise ValueError("Invalid SSH destination format. Expected user@host:/path")

        # Split host part and path part
        host_part, self.remote_path = dest.split(":", 1)

        # Extract username and hostname
        if "@" in host_part:
            self.username, self.hostname = host_part.split("@", 1)
        else:
            raise ValueError("Invalid SSH destination format. Expected user@host:/path")

        # Get password
        self.password = os.environ.get("SSH_PASSWORD")
        if not self.password:
            Logger.warning(
                "SSH_PASSWORD environment variable not set. Trying key-based authentication."
            )

    def connect(self):
        """Establish SSH connection"""
        try:
            # Create SSH client
            self.ssh_client = SSHClient()
            self.ssh_client.set_missing_host_key_policy(AutoAddPolicy())

            # Connect to SSH server
            if self.ssh_key_file:
                self._connect_with_keys()
            else:
                self.ssh_client.connect(
                    hostname=self.hostname,
                    username=self.username,
                    password=self.password,
                )
            self.sftp = self.ssh_client.open_sftp()
            self.scp_client = SCPClient(self.ssh_client.get_transport())
            Logger.success(f"Connected to {self.username}@{self.hostname}")

        except Exception as e:
            Logger.error(f"Failed to connect to SSH server: {e}")
            self.close()

    def _connect_with_keys(self):
        """Connect using SSH keys"""
        connected = False

        if os.path.exists(self.ssh_key_file):
            try:
                # Determine key class based on file extension
                key_class = None
                if "ed25519" in self.ssh_key_file:
                    key_class = paramiko.Ed25519Key
                elif "ecdsa" in self.ssh_key_file:
                    key_class = paramiko.ECDSAKey
                elif "dsa" in self.ssh_key_file:
                    key_class = paramiko.DSSKey
                else:
                    key_class = paramiko.RSAKey

                # Try to load and use the key
                private_key = key_class.from_private_key_file(self.ssh_key_file)
                self.ssh_client.connect(
                    hostname=self.hostname, username=self.username, pkey=private_key
                )
                Logger.success(f"Successfully connected using key: {self.ssh_key_file}")
                connected = True
            except Exception as e:
                Logger.warning(f"Failed to connect using key {self.ssh_key_file}: {e}")

        # If no key files work, try connecting without explicit key (may use SSH agent)
        if not connected:
            try:
                self.ssh_client.connect(hostname=self.hostname, username=self.username)
                Logger.success("Connected using SSH agent or other method")
            except Exception as e:
                Logger.error(f"Failed to connect using SSH agent or other methods: {e}")
                raise

    def create_remote_dirs(self):
        """Recursively create remote directories"""
        try:
            # Recursively create directories
            dirs = self.remote_path.strip("/").split("/")
            current_dir = "/"
            for dir_part in dirs:
                if dir_part:
                    current_dir = f"{current_dir}{dir_part}/"
                    try:
                        self.sftp.stat(current_dir)
                    except FileNotFoundError:
                        self.sftp.mkdir(current_dir)

            Logger.success(f"Remote directories created at {self.remote_path}")

        except Exception as e:
            Logger.error(f"Failed to create remote directories: {e}")

    def copy_files(self, source_path):
        """Copy files to remote target"""

        try:
            # Use SCP to copy files
            if os.path.isdir(source_path):
                for item in os.listdir(source_path):
                    local_item_path = os.path.join(source_path, item)
                    remote_item_path = os.path.join(self.remote_path, item).replace(
                        "\\", "/"
                    )
                    self.scp_client.put(local_item_path, remote_item_path, recursive=True)
            else:
                # If source_path is a file
                self.scp_client.put(source_path, self.remote_path)

            Logger.success("Files copied to SSH destination successfully.")

        except Exception as e:
            Logger.error(f"Failed to copy files to SSH destination: {e}")

    def close(self):
        """Close SSH connection"""
        if self.scp_client:
            self.scp_client.close()
        if self.sftp:
            self.sftp.close()
        if self.ssh_client:
            self.ssh_client.close()
        Logger.info("SSH connection closed.")
